import re
from typing import List
from langchain.prompts import SystemMessagePromptTemplate, ChatPromptTemplate
from langchain_experimental.plan_and_execute import (
    PlanAndExecute,
    load_agent_executor,
    load_chat_planner,
)
from langchain_openai import ChatOpenAI
from langchain.agents import tool


llm = ChatOpenAI(
    openai_api_key="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",
    openai_api_base="https://openrouter.ai/api/v1",
    model_name="deepseek/deepseek-chat-v3-0324:free",
)


plan_tmp = """\
    首先请你理解这个问题，然后设计解决这个问题的计划，设计计划过程中必须遵守以下规则：
    1、计划的表头必须使用“plan:”开始，然后是步骤的编号列表
    2、以最少的计划步骤准确的完成计划设计
    3、设计计划时优先考虑使用工具列表中的工具完成计划任务，工具列表：{Tools}
    4、设计计划时需考虑计划的依赖性和连贯性，尽量使用原始问题中的相关信息调用对应工具完成计划任务
    5、调用工具时严格遵守工具的参数要求，具体信息参考工具列表详细描述
    6、计划步骤中严禁出现重复计划或者相似的重复计划
    7、计划的最后一步需考虑之前的计划是否能解决原始问题
    8、在计划的末尾必须使用“<END_OF_PLAN>”表示计划结束
    9、相关工具返回结果中“<END_RUN>”表示结束当前全部计划
    
    以下为提供设计计划的工具列表详细描述：
    {Tools_Describe}
"""


@tool()
def stylesheet(*args, **kwargs) -> List[str]:
    "网页开发所需的CSS框架"
    return ["<bootstrap@5.3.0>", "<TailwindCSS@3.3.3>"]


@tool()
def bingtu(*args, **kwargs) -> List[str]:
    "这个工具可以提供使用哪些包生成饼图样式"

    return ["使用<chart.js@4.4.8>这个包，并使用'doughnut'这个类型"]


@tool()
def zhuzhuang(*args, **kwargs) -> List[str]:
    "这个工具可以提供使用哪些包生成柱状图样式"

    return ["使用<chart.js@4.4.8>这个包，并使用'bar'这个类型"]


@tool()
def zhexian(*args, **kwargs) -> List[str]:
    "这个工具可以提供使用哪些包生成折线图样式"

    return ["使用<chart.js@4.4.8>这个包，并使用'line'这个类型"]


@tool()
def liebiao(*args, **kwargs) -> List[str]:
    "这个工具可以提供使用哪些包生成折线图样式"

    return ["使用<chart.js@4.4.8>这个包，并使用'list'这个类型"]


tools = [stylesheet, bingtu, zhuzhuang, zhexian, liebiao]
tool_strings = []
for t in tools:
    args_schema = re.sub("}", "}}", re.sub("{", "{{", str(t.args)))
    tool_strings.append(f"{t.name}: {t.description}, args: {args_schema}")
    formatted_tools = "\n".join(tool_strings)
    tool_names = ", ".join([t.name for t in tools])

sys_prompt_tmp = SystemMessagePromptTemplate.from_template(plan_tmp)
chat_sys_tmp = ChatPromptTemplate.from_messages([sys_prompt_tmp])
sys_prompt = chat_sys_tmp.format(Tools=tool_names, Tools_Describe=formatted_tools)

planner = load_chat_planner(llm, system_prompt=sys_prompt)
executor = load_agent_executor(llm, tools, verbose=True)


summary = """\
    ### 网络安全摘要信息

#### 资产网络攻击统计相关数据
- **资产ID**: `0fd4b9d96d6f5ecab6c6a8b6809f5ae3`
- **资产名称**: `taomingwei-node3-ubuntu`
- **总事件数量**: 13个
- **告警时间范围**：
  - 第一条告警时间：2022年4月3日 09:52:39
  - 最新告警时间：2025年4月8日 09:12:25

#### 资产网络攻击类型分布
- **攻击ID及描述**：
  - T1564.001：通过命令行创建隐藏文件和目录，告警级别为3。
  - T1078.003：通过 Grep 命令查找系统上包含 Password 关键字的文件，告警级别为3。
  - T1082：可疑 Modprobe 文件事件，告警级别为4。
  - T1059：
    - 检测到Linux tmp目录程序运行和检测到Linux可疑shell命令，告警级别分别为4（第3、5条数据）。
    - Linux 受限 Shell 通过 Linux 二进制突破，告警级别为3。
    - 检测到Linux perl反弹shell2，威胁级别: 4
  - T1027.001 (填充二进制文件)
  - T1070.002 (清除Linux系统日志)

- **攻击阶段**：
  - `TA0007`（第2条数据）
  - `TA0002`（第3、4、5条数据，以及检测到Linux perl反弹shell事件）
  - `TA0004`（填充二进制文件）
  - `TA0005`（清除Linux系统日志）

#### 资产漏洞相关情况
- **检测方式**：
  - 隐藏文件目录检测：通过命令行创建隐藏文件和目录。
  - 敏感命令执行检测：使用 Grep 命令查找系统上包含 Password 关键字的文件。
  - 发现系统信息检测：检测 modprobe 配置文件中涉及内核模块的文件事件。
  - 检测到Linux tmp目录程序运行（第3条数据）
  - Linux 受限 Shell 通过 Linux 二进制突破（第4条数据）
  - 检测到Linux可疑shell命令（第5条数据）
  - 填充二进制文件: 检测通过 DD 指令填充二进制的行为。
  - 清除Linux系统日志: 黑客通常在入侵结束后对入侵痕迹进行清理，例如对所有可能记录其攻击行为的日志进行删除操作。

#### 资产风险相关情况
- **告警级别**：
  - 第3、5条数据：告警级别为4。
  - 第4条数据：告警级别为3。
  - 填充二进制文件: 级别3 (第6条)
  - 清除Linux系统日志: 级别3 (第7条)
  - 检测到Linux 可疑命令: 级别4 (第8条)

#### 资产安全防护状态
- **修复建议**：
  - 排查和结束相关可疑进程。
  - 排查和删除相关文件。
  - 排查和删除相关系统设置。
  - 排查和删除相关注册表项。
  - 填充二进制文件:
    1. 通过任务管理排查和结束相关可疑进程。
    2. 通过资源管理器排查和删除相关文件。
    3. 排查和删除相关可疑系统设置。
    4. 排查和删除相关可疑注册表项。
  - 清除Linux系统日志:
    1. 确认删除的日志文件是否可疑，排除正常的运维与文件更新操作。
    2. 确认操作可疑的情况下，进一步排查导致文件删除的进程名以及对应的用户名，必要时结束进程。
    3. 对可疑进程的后续操作进行排查。
  - 检测到Linux 可疑命令:
    排查命令上下文分析命令是否恶意，若为恶意及时结束相关进程。

### 总结
资产 `taomingwei-node3-ubuntu` 遭遇了多次网络安全事件：
1. 检测到Linux可疑命令。
2. 检测到Linux perl反弹shell。

建议立即排查相关网络连接并采取措施防止进一步的恶意活动。
"""
try:
    agent = PlanAndExecute(planner=planner, executor=executor, verbose=True)
    input = """\
        根据<摘要信息>生成一个html网页代码，尽量使用工具提供的第三方库，如果没有对应的库可以使用则自由发挥，如下是计划步骤：
        1、理解<摘要信息>内容，并进行分析哪些数据需要使用工具进行渲染
        2、根据<摘要信息>生成一个html网页框架
        3、根据工具选择第三方库生成html数据样式
        4、最后提供完整的HTML代码模板

        <摘要信息>：
        {summary}
    """
    input = input.format(summary=summary)
    ret = agent.invoke(input, verbose=True)
    print(ret)
except Exception as e:
    print("doc summary error:{}, input:{}", e, input)
